# 1. 导入所需库
import pandas as pd  # 处理Excel数据
import matplotlib.pyplot as plt  # 绘图
from sklearn.tree import DecisionTreeClassifier, plot_tree  # 决策树模型和可视化
from sklearn.model_selection import train_test_split  # 划分训练集和测试集
from sklearn.metrics import accuracy_score  # 评估准确率
import joblib
import areaposition
import matplotlib.pyplot as plt
import numpy as np  # 确保顶部导入了 numpy

plt.rcParams["font.family"] = ["SimHei"]  # 支持中文


def train():
    # 2. 读取数据（替换为你的Excel文件路径）
    # 假设Excel中"丢弃的牌"列：若丢弃新牌则记为0，若丢弃旧牌则记为对应牌值（如9）
    data = pd.read_excel("mahjong_data.xlsx")  # 请将此处改为你的文件路径

    # 3. 数据预处理（提取特征X和标签y）
    # 特征X：13张旧牌 + 1张新牌（共16个特征）
    x = data[["card01", "card02", "card03", "card04", "card05",
              "card06", "card07", "card08", "card09", "card10",
              "card11", "card12", "card13", "card14"]]

    # 标签y：最终丢弃的牌（目标变量）
    y = data["dis"]

    # 4. 划分训练集（80%）和测试集（20%）
    x_train, x_test, y_train, y_test = train_test_split(
        x, y, test_size=0.2, random_state=42  # random_state固定划分方式，保证结果可复现
    )

    # 5. 训练决策树模型（CART分类树，适合多类别决策）
    model = DecisionTreeClassifier(
        criterion="gini",  # CART用基尼系数
        max_depth=5,  # 限制树深度，避免过拟合（可根据数据量调整，5-10为宜）
        min_samples_leaf=3  # 叶子节点最少3个样本，避免学细节
    )
    model.fit(x_train, y_train)  # 用训练集训练模型

    # 6. 评估模型性能
    y_pred = model.predict(x_test)  # 用测试集预测
    accuracy = accuracy_score(y_test, y_pred)  # 计算准确率
    print(f"模型准确率：{accuracy:.2f}（越高越好，1.0为完美）")

    # 7. 可视化决策树（查看模型学到的麻将取舍规则）
    plt.figure(figsize=(20, 10))  # 画布大小
    plot_tree(
        model,
        feature_names=x.columns,  # 特征名称（13旧牌+1新牌）
        class_names=[str(cls) for cls in model.classes_],  # 标签名称（丢弃的牌值）
        filled=True,  # 用颜色深浅表示节点纯度
        rounded=True,  # 圆角矩形
        fontsize=8  # 字体大小
    )
    plt.title("麻将取舍决策树规则")
    plt.show()  # 显示决策树

    joblib.dump(model, "mahjong_model.pkl")  # 保存为.pkl文件
    print("模型已保存：mahjong_model.pkl")


def predicted():
    # 8. 用新数据预测（可选：输入一组新的13+1张牌，看模型推荐丢哪张）
    # 示例：假设新牌组为[1,2,3,4,5,6,7,8,9,1,2,3,4,5]（前13为旧牌，最后1为新牌）
    model = joblib.load("mahjong_model.pkl")  # 路径与保存时一致
    print("模型加载成功！")

    new_data = [item.name for item in areaposition.card_position][1:]

    # 假设训练时的特征列名为 feature_names（例如：['f1', 'f2', ..., 'fn']）
    feature_names = ["card01", "card02", "card03", "card04", "card05",
                     "card06", "card07", "card08", "card09", "card10",
                     "card11", "card12", "card13", "card14"]  # 替换为你训练时的实际特征名

    # 将 new_data 转为 DataFrame（保持与训练时相同的列名）
    new_data_df = pd.DataFrame([new_data], columns=feature_names)

    # 用 DataFrame 预测（带特征名，匹配训练时的格式）
    predicted_drop = model.predict(new_data_df)

    print(f"new_data:{new_data}")
    # new_data = [[1, 2, 3, 4, 5, 6, 7, 8, 9, 1, 2, 3, 4, 5]]
    # predicted_drop = model.predict(np.array(new_data).reshape(1, -1))  # 先转数组再reshape
    # predicted_drop = model.predict(new_data)
    print(f"对于新牌组，模型推荐丢弃：{predicted_drop[0]}")
    return predicted_drop


if __name__ == "__main__":
    train()
